from TLSMapper.TLSSUT import *
from TLSMapper.TLSProtocol import *
from datetime import datetime
from aalpy.utils import visualize_automaton
from aalpy.learning_algs import run_Lstar
from aalpy.oracles import RandomWMethodEqOracle, StatePrefixEqOracle
import os
# 'OpenSSL', 'WolfSSL', 'MatrixSSL', 'GnuTLS', 'mbedTLS', 'BoringSSL'

ciphersuites=[
              CipherSuite.TLS_AES_128_GCM_SHA256,
            #   CipherSuite.TLS_AES_256_GCM_SHA384,
            #   CipherSuite.TLS_ECDHE_ECDSA_WITH_AES_256_GCM_SHA384,
            #   CipherSuite.TLS_DHE_RSA_WITH_CHACHA20_POLY1305_SHA256,
            #   CipherSuite.TLS_ECDHE_RSA_WITH_AES_256_GCM_SHA384
              ]
# # ciphersuites=[0x1302,0x1303,0x1301,0x00ff]
ciphersuites=[0xc02c,0xc030,0x00a3,0x009f,0xcca9,0xcca8,0xccaa,0xc02b,0xc02f,0x00a2,0x009e,0xc0ac,0xc0ad,0x006b,0x006a,0xc023,0xc027,0x0067,0x0040,0xc00a,0xc014,0x0039,0x0038,0xc009,0xc013,0x0033,0x0032,0x009d,0x009c,0x003d,0x003c,0x0035,0x002f,0x00ff]

# ciphersuites=[0x1302,0x1303,0x1301,0xc02c,0xc030,0x009f,0xcca9,0xcca8,0xccaa,0xc02b,0xc02f,0x009e,0xc024,0xc028,0x006b,0xc023,0xc027,0x0067,0xc00a,0xc014,0x0039,0xc009,0xc013,0x0033,0x009d,0x009c,0x003d,0x003c,0x0035,0x002f,0x00ff]

'''
###############################################
####                TLS 1.2                ####
###############################################
'''
# alphabet = ['ClientHello','Certificate', 'ClientKeyExchange','CertificateVerify','ChangeCipherSpec','Finish', 'ApplicationData', 'EmptyCertificate', 'TLS12ReClientHello']
# tlsimp = TLSProtocol(version=(3,3), cipher_suite=ciphersuites, implementation = 'mbedTLS', verify = True, fuzz_tls=(False, None, None))
# tlsimp = TLSProtocol(version=(3,3), cipher_suite=ciphersuites, implementation = 'WolfSSL', verify = True, fuzz_tls=(False, None, None))

# tlsimp = TLSProtocol(version=(3,3), cipher_suite=ciphersuites, implementation = 'MatrixSSL', verify = True, fuzz_tls=(False, None, None))



'''
###############################################
####                TLS 1.3                ####
###############################################
'''
'''
select alphabet
'''
# alphabet = ['ClientHelloEmptyKeyShare','ClientHello', 'ChangeCipherSpec', 'Certificate',  'CertificateVerify', 'Finish', 'KeyUpdate', 'ApplicationData','ResumptionClientHello']
# # alphabet = ['ClientHelloEmptyKeyShare','ClientHello', 'ChangeCipherSpec', 'Certificate',  'CertificateVerify', 'Finish', 'ApplicationData', 'ResumptionClientHello', 'EmptyCertificate']
# # alphabet = ['ClientHelloEmptyKeyShare','ClientHello', 'ChangeCipherSpec', 'Certificate',  'CertificateVerify', 'Finish', 'ApplicationData', 'ResumptionClientHello', 'EmptyCertificate']
# # alphabet = ['ClientHelloEmptyKeyShare','ClientHello', 'ChangeCipherSpec', 'Certificate',  'CertificateVerify', 'Finish', 'ApplicationData', 'KeyUpdate']
# # alphabet = ['ClientHello']
# # alphabet = ['ClientHelloEmptyKeyShare','ClientHello', 'ChangeCipherSpec', 'Certificate',  'CertificateVerify', 'Finish', 'ApplicationData','ResumptionClientHello']
# alphabet = ['ClientHelloEmptyKeyShare','ClientHello', 'ChangeCipherSpec', 'Certificate',  'CertificateVerify', 'Finish', 'KeyUpdate', 'ApplicationData','ResumptionClientHello']
# alphabet = ['ClientHelloEmptyKeyShare','ClientHelloTLS13', 'ChangeCipherSpec', 'Certificate',  'CertificateVerify', 'Finish', 'ApplicationData','ClientHelloTLS12','ResumptionClientHello']
# alphabet = ['ClientHelloEmptyKeyShare','ClientHelloTLS12', 'ChangeCipherSpec', 'Certificate',  'CertificateVerify', 'Finish', 'ApplicationData','ResumptionClientHello']
# alphabet = ['ClientHelloEmptyKeyShare','ClientHello', 'ChangeCipherSpec', 'Certificate',  'CertificateVerify', 'Finish', 'ApplicationData','ClientHelloTLS12','ResumptionClientHello']
# alphabet = ['ClientHelloEmptyKeyShare','ClientHello', 'ChangeCipherSpec', 'Certificate',  'CertificateVerify', 'Finish', 'ApplicationData',  'CertificateVerify','ResumptionClientHello']
# alphabet = ['ClientHelloEmptyKeyShare','ClientHello', 'ChangeCipherSpec','ClientHelloTLS12']

# alphabet = ['ClientHelloEmptyKeyShare','ClientHello', 'ChangeCipherSpec', 'Certificate',  'CertificateVerify', 'Finish', 'ApplicationData','ClientHelloTLS12','ResumptionClientHello']

# alphabet = ['ClientHelloTLS12', 'ChangeCipherSpec']

# alphabet = ['ClientHelloEmptyKeyShare','ClientHello', 'ChangeCipherSpec', 'Certificate',  'CertificateVerify', 'Finish', 'ApplicationData', 'ResumptionClientHello', 'Finish', 'KeyUpdate', 'ApplicationData', 'ApplicationData','ClientHello','ClientHello','ClientHello', 'Finish']
# alphabet = ['ClientHello', 'Certificate', 'CertificateVerify', 'Finish', 'ResumptionClientHello', 'ResumptionClientHello', 'ResumptionClientHello']
# alphabet1 = ['ClientHello', 'ChangeCipherSpec', 'Finish', 'ResumptionClientHelloAP', 'EndOfEarlyData', 'Finish']
# alphabet1 = ['ClientHello', 'ChangeCipherSpec', 'Finish', 'ApplicationData']
# alphabet = ['ClientHello', 'ChangeCipherSpec', 'Finish', 'ResumptionClientHelloAP', 'EndOfEarlyData', 'Finish']
# alphabet1 = ['ClientHello', 'ChangeCipherSpec', 'Certificate', 'CertificateVerify', 'Finish', 'ResumptionClientHelloAP', 'Finish']
# alphabet2 = ['ClientHello', 'ChangeCipherSpec', 'Certificate', 'CertificateVerify', 'Finish', 'ResumptionClientHello', 'Finish']
alphabet = ['ClientHello', 'ClientKeyExchange','ChangeCipherSpec', 'Finish']
encrypted_messages = [
    'Finish',                    # 第一个加密的握手消息
    'ApplicationData',           # 所有应用数据
    'ClosureAlert',              # 加密关闭
    'ErrorAlert',                # 通常在加密后发送
    'KeyUpdate',                 # TLS 1.3 密钥更新
    'EndOfEarlyData',            # TLS 1.3 0-RTT 结束
    'ResumptionClientHelloAP',   # 包含加密的 ApplicationData 和 EndOfEarlyData
]




'''
select tlsimp
'''
# tlsimp = TLSProtocol(version=(3,4), cipher_suite=ciphersuites, implementation = 'mbedTLS', verify = True, fuzz_tls=(False, None, None))
# tlsimp = TLSProtocol(version=(3,4), cipher_suite=ciphersuites, implementation = 'mbedTLS', verify = True, fuzz_tls=(False, None, None))
# tlsimp = TLSProtocol(version=(3,4), cipher_suite=ciphersuites, implementation = 'WolfSSL', verify = True, fuzz_tls=(False, None, None))
# tlsimp = TLSProtocol(version=(3,4), cipher_suite=ciphersuites, implementation = 'BoringSSL', verify = True, target=('127.0.0.1',35001), fuzz_tls=(False, None, None))
# tlsimp = TLSProtocol(version=(3,4), cipher_suite=ciphersuites, implementation = 'MatrixSSL', verify = True, fuzz_tls=(False, None, None))
# tlsimp = TLSProtocol(version=(3,4), cipher_suite=ciphersuites, implementation = 'MatrixSSL', verify = True, fuzz_tls=(False, None, None))



# os.environ['LD_LIBRARY_PATH'] = '/home/kai/Desktop/tls_infer/openssl-gcov' + os.environ.get('LD_LIBRARY_PATH', '')
# print(os.environ['LD_LIBRARY_PATH'])
# tlsimp = TLSProtocol(version=(3,4), cipher_suite=ciphersuites, implementation = 'OpenSSL', verify = True, target=('127.0.0.1',44331), fuzz_tls=(False, None, None),imp_ver='1.1.1m')
# tlsimp = TLSProtocol(version=(3,4), cipher_suite=ciphersuites, implementation = 'OpenSSL', verify = False, target=('127.0.0.1',44330), fuzz_tls=(False, None, None),imp_ver='1.1.1m')
tlsimp = TLSProtocol(version=(3,3), cipher_suite=ciphersuites, implementation = 'OpenSSL', verify = False, target=('127.0.0.1',12345), fuzz_tls=(False, None, None))
# tlsimp = TLSProtocol(version=(3,4), cipher_suite=ciphersuites, implementation = 'OpenSSL', verify = True, target=('127.0.0.1',44334), fuzz_tls=(False, None, None))

# for no enc

# tlsimp.not_enc = True

'''
select impeletation
'''
sul = TLSSUT( TLSpro = tlsimp )


# sul = TLSSUT(keyfile='./certs/client-key.pem', certfile='./certs/client-cert.pem', TLSpro = tlsimp)
# sul = TLSSUT(keyfile='./testkeys/RSA/2048_RSA_KEY.pem', certfile='./testkeys/RSA/2048_RSA.pem', TLSpro = tlsimp)


# sul.target_cmd = ['/home/kai/Desktop/server/openssl-1.1.1w/apps/openssl', 's_server', '-key', './key/deserver.key', '-cert', './key/server.cer', '-CAfile', './key/ca.cer', '-accept', '4433', '-verify', '1', '-HTTP']
# sul.target_cmd = ['/home/kai/Desktop/server/openssl-1.1.1w-patch/apps/openssl', 's_server', '-key', './key/deserver.key', '-cert', './key/server.cer', '-CAfile', './key/ca.cer', '-accept', '4433', '-verify', '1', '-HTTP', '-naccept', '1']
# sul.target_cmd = ['/home/kai/Desktop/tls_infer/openssl-gcov/apps/openssl', 's_server', '-key', './key/deserver.key', '-cert', './key/server.cer', '-CAfile', './key/ca.cer', '-accept', '44330', '-verify', '1', '-HTTP', '-naccept', '1']
# sul.target_cmd = ['/home/kai/Desktop/tls_infer/openssl-gcov/apps/openssl', 's_server', '-key', './key/deserver.key', '-cert', './key/server.cer', '-CAfile', './key/ca.cer', '-accept', '44330', '-HTTP', '-naccept', '1']
# sul.target_cmd = ['/home/kai/Desktop/tls_infer/openssl-gcov/apps/openssl', 's_server', '-key', './key/deserver.key', '-cert', './key/server.cer', '-CAfile', './key/ca.cer', '-accept', '44331', '-verify', '1', '-HTTP', '-naccept', '2']
# sul.target_cmd = ['/home/kai/Desktop/tls_infer/openssl-gcov/apps/openssl', 's_server', '-key', './key/deserver.key', '-cert', './key/server.cer', '-CAfile', './key/ca.cer', '-accept', '44331', '-verify', '1', '-HTTP', '-naccept', '2']
# sul.target_cmd = ['openssl', 's_server', '-key', './key/deserver.key', '-cert', './key/server.cer', '-accept', '44334', '-early_data', '-naccept', '2', '-HTTP']
# sul.target_cmd = ['/home/kai/Desktop/tls_infer/openssl-gcov/apps/openssl', 's_server', '-key', './key/deserver.key', '-cert', './key/server.cer', '-CAfile', './key/ca.cer', '-accept', '44333', '-verify', '1', '-HTTP']


# sul.target_ip = '127.0.0.1'
# sul.target_port = 44332
# alphabet1 = ['ClientHello', 'ChangeCipherSpec', 'Finish', 'ApplicationData']
alphabet1 = ['ClientHello', 'ClientKeyExchange','ChangeCipherSpec', 'Finish']


sul.query(alphabet1)
# sul.query(alphabet2)



# eq_oracle = RandomWMethodEqOracle(alphabet, sul)
# # # eq_oracle = StatePrefixEqOracle(alphabet, sul,walks_per_state=20, walk_len=10)
# # eq_oracle = StatePrefixEqOracle(alphabet, sul,walks_per_state=5, walk_len=5)

# imp = tlsimp.implementation
# ver = tlsimp.imp_ver
# if tlsimp.version == (3,4):
#     tv = 'tls1.3'
# else:
#     tv = 'tls1.2'
# folder_name = f"{imp}_{ver}_{tv}"
# os.makedirs(folder_name, exist_ok=True)
# save_dir = f"./{imp}_{ver}_{tv}/model.dot"
# # print(save_dir)

# model = run_Lstar(alphabet, sul, eq_oracle, 'mealy', cache_and_non_det_check=False,debug_save=folder_name)
# model.save(file_path=save_dir)
# visualize_automaton(model)

# os.system("./cov_script")
